import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
df = pd.read_excel('stock.xlsx', dtype={'code': 'str'})
df.set_index('code', inplace=True)
print('显示代码为002522的股票基本面数据')
a = df.loc['002522']
print(a)
print('行业数')
b = len(df.industry.unique())
print(b)
print('地区数')
c = len(df.area.unique())
print(c)

d = df.groupby('area').area.count().sort_values(ascending=False)
print(d)
year = df.timeToMarket.astype('str').str[:4]  # 转换为字符串，提取前4位的年份
yearnum = df.groupby(year).name.count()  # 按年份统计，得到每年股票发行量
print(yearnum)
plt.rcParams['font.sans-serif']=['SimHei']
plt.rcParams['axes.unicode_minus']=False
yearnum[yearnum.index != '0'].plot(fontsize=14, title='年IPO数量') # 数据集中有几只股票没有发行年份（年份为0）, 作图时排除0年份
plt.savefig("年IPO数量.png")
plt.show()

# 市盈率的计算
x = df.pe.mean()       		# 简单的算术平均pe
print(x)
y = df[df.pe > 0].pe.mean() 	# 剔除亏损股票后计算pe均值
print(y)
df['tvalue'] = 4 * df.esp * df.pe * df.totals  	# 计算总市值，增加新列tvalue
z = np.sum(df.pe * df.tvalue) / df.tvalue.sum()   # 计算以市值为权重的加权pe
print(z)
df['board'] = df.index.str[:2]     # 取code的前2个字符，新增board列
# 按板块类型统计pe均值，计数
n = df.groupby('board').pe.agg([('pe均值', 'mean'), ('股票数', 'count')])
print(n)
